QteaTensor via jax
- qredtea.jaxapi.default_jax_backend(device='cpu', dtype=jax.numpy.complex64)[source]
Generate a default tensor backend for symmetric tensors, i.e., with a
QteaJaxTensor.Arguments
- dtypedata type, optional
Data type for JAX. Default to jax.complex64
- devicedevice specification, optional
Default to “cpu”. Available: “cpu”, “gpu”, “xla”
Returns
tensor_backend :
TensorBackend
- qredtea.jaxapi.default_abelian_jax_backend(device='cpu', dtype=jax.numpy.complex128)[source]
Provide a backend using jax and Abelian tensors.
Arguments
- dtypedata type, optional
Data type for JAX. Default to jax.complex64
- devicedevice specification, optional
Default to “cpu”. Available: “cpu”, “gpu”, “xla”
Returns
tensor_backend :
TensorBackend
Tensor class based on JAX; JAX supports both CPU and GPU in one framework.
- class qredtea.jaxapi.QteaJaxTensor(links, ctrl='Z', are_links_outgoing=None, base_tensor_cls=None, dtype=jax.numpy.complex64, device=None, requires_grad=None)[source]
Tensor for Quantum TEA based on JAX tensors.
- add_update(other, factor_this=None, factor_other=None)[source]
Inplace addition as self = factor_this * self + factor_other * other.
Arguments
- othersame instance as self
Will be added to self. Unmodified on exit.
- factor_thisscalar
Scalar weight for tensor self.
- factor_otherscalar
Scalar weight for tensor other
- property are_links_outgoing
Define property of outgoing links as property (always False).
- attach_dummy_link(position, is_outgoing=True)[source]
Attach dummy link at given position (inplace update).
- concatenate_vectors(vectors, dtype, dim=None)[source]
Concatenate vectors of the underlying numpy / cupy / torch / etc tensors.
Arguments*
- vectorslist
List of one-dimensional arrays.
- dtypedata type
Data type of concatenated vectors.
- dimint | None
Total dimension of concatenated vectors. If None, calculated on the fly. Default to None
Details
Used to concatenate singular values for symmetric tensors in SVD, which is needed as jax and tensorflow do not support x[:] assignments.
- convert(dtype=None, device=None, stream=None)[source]
Convert underlying array to the specified data type inplace.
- static convert_operator_dict(op_dict, params=None, symmetries=None, generators=None, base_tensor_cls=None, dtype=jax.numpy.complex64, device='cpu')[source]
Iterate through an operator dict and convert the entries. Converts as well to rank-4 tensors.
Arguments
- op_dictinstance of
TNOperators Contains the operators as xp.ndarray.
- paramsdict, optional
To resolve operators being passed as callable.
- symmetries: list, optional, for compatability with symmetric tensors.
Must be empty list.
- generatorslist, optional, for compatability with symmetric tensors.
Must be empty list.
- base_tensor_clsNone, optional, for compatability with symmetric tensors.
No checks on this one here.
- dtypedata type for xp, optional
Specify data type. Default to jnp.complex64
- devicestr
Device for the simulation. Available “cpu” and “gpu” Default to “cpu”
Details
The conversion to rank-4 tensors is useful for future implementations, either to support adding interactions with a bond dimension greater than one between them or for symmetries. We add dummy links of dimension one. The order is (dummy link to the left, old link-1, old link-2, dummy link to the right).
- op_dictinstance of
- property device
Device where the tensor is stored.
- property dtype
Data type of the underlying arrays.
- property dtype_eps
Data type’s machine precision.
- eig()[source]
Compute eigenvalues and eigenvectors of a two-leg tensor
Return
- eigvals, eigvecsinstances of
QteaJaxTensor Eigenvalues and corresponding eigenvectors of input tensor.
- eigvals, eigvecsinstances of
- eig_api(matvec_func, links, conv_params, args_func=None, kwargs_func=None)[source]
Interface to hermitian eigenproblem
Arguments
- matvec_funccallable
Mulitplies “matrix” with “vector”
- linkslinks according to
QteaTensor Contain the dimension of the problem.
- conv_paramsinstance of
TNConvergenceParameters Settings for eigenproblem with Arnoldi method.
args_func : arguments for matvec_func
kwargs_func : keyword arguments for matvec_func
Returns
eigenvalues : scalar
eigenvectors : instance of
QteaTensor
- eig_api_arpack(matvec_func, links, conv_params, args_func=None, kwargs_func=None)[source]
Interface to hermitian eigenproblem via Arpack.
- eig_api_qtea(matvec_func, conv_params, args_func=None, kwargs_func=None)[source]
Interface to hermitian eigenproblem via qtealeaves.solvers. Arguments see eig_api.
- eigh()[source]
Compute eigenvalues and eigenvectors of a two-leg hermitian tensor.
Return
- eigvals, eigvecsinstances of
QteaJaxTensor Eigenvalues and corresponding eigenvectors of input tensor.
- eigvals, eigvecsinstances of
- eigvalsh()[source]
Compute eigenvalues of a two-leg Hermitian tensor
Return
- eigvalsjax.Array
Eigenvalues of input tensor.
- einsum(einsum_str, *others)[source]
Call to einsum with self as first tensor.
Arguments
- einsum_strstr
Einsum contraction rule.
- other: List[
QteaJaxTensors] 2nd, 3rd, …, n-th tensor in einsum rule as positional arguments.
Results
- tensor
QteaJaxTensor Contracted tensor according to the einsum rules.
Details
The call
np.einsum(einsum_str, x.elem, y.elem, z.elem)translates intox.einsum(einsum_str, y, z)for x, y, and z beingQteaJaxTensor.
- property elem
Elements of the tensor.
- elementwise_abs_smaller_than(value)[source]
Return boolean if each tensor element is smaller than value
- expm(fuse_point=None, prefactor=1)[source]
Take the matrix exponential with a scalar prefactor, i.e., Exp(prefactor * self).
Parameters
- fuse_pointint, optional
If given, reshapes the tensor into a matrix by fusing links up to INCLUDING fuse_point into one, and links after into the second dimension. To compute the exponential of a 4-leg tensor, for example, by fusing (0,1),(2,3), set fuse_point=1. Default is None.
- prefactorfloat, optional
Prefactor of the tensor to be exponentiated. Default to 1.
Return
- matinstance of
QteaTensor Exponential of input tensor.
- eye_like(link)[source]
Generate identity matrix.
Arguments
- selfinstance of
QteaTensor Extract data type etc from this one here.
- linksame as returned by links property, here integer.
Dimension of the square, identity matrix.
- selfinstance of
- classmethod from_elem_array(tensor, dtype=None, device=None)[source]
New QteaJaxTensor from array
Arguments
- tensorjax.Array
Array for new tensor.
- dtypedata type, optional
Can allow to specify data type. If not None, it will convert. Default to None
- classmethod from_qteatensor(qteatensor, dtype=None, device=None)[source]
Convert QteaTensor based on numpy/cupy into QteaJaxTensor.
- fuse_links_update(fuse_low, fuse_high, is_link_outgoing=True)[source]
Fuses one set of links to a single link (inplace-update).
Parameters
- fuse_lowint
First index to fuse
- fuse_highint
Last index to fuse.
Example: if you want to fuse links 1, 2, and 3, fuse_low=1, fuse_high=3. Therefore the function requires links to be already sorted before in the correct order.
- get_submatrix(row_range, col_range)[source]
Extract a submatrix of a rank-2 tensor for the given rows / cols.
- static is_cpu_static(device_str)[source]
Check if device is CPU or not.
Parameters
- device_strstr
Check for this string if it corresponds to a CPU.
Returns
- is_cpubool
True if device is a CPU.
- static is_gpu_static(device_str)[source]
Check if device is GPU or not.
Parameters
- device_strstr
Check for this string if it is a GPU device.
Returns
- is_gpubool
True if device is a GPU.
- is_implemented_device(query)[source]
Check if argument query is an implemented device.
Parameters
- querystr
String to be tested if it corresponds to a device implemented with this tensor.
Returns
- is_implementedbool
True if string is available as device.
- is_xla(query=None)[source]
Check if device is XLA or not.
Parameters
- querystr | None, optional
If given, check for this string. If None, self.device will be checked. Default to None.
Returns
- is_xlabool
True if device is a XLA.
- static is_xla_static(device_str)[source]
Check if device is XLA or not.
Parameters
- device_strstr
Check for this string if it is a XLA device.
Returns
- is_xlabool
True if device is a XLA.
- kron(other, idxs=None)[source]
Perform the kronecker product between two tensors. By default, do it over all the legs, but you can also specify which legs should be kroned over. The legs over which the kron is not done should have the same dimension.
Parameters
- otherQteaTensor
Tensor to kron with self
- idxsTuple[int], optional
Indexes over which to perform the kron. If None, kron over all indeces. Default to None.
Returns
- QteaTensor
The kronned tensor
Details
Performing the kronecker product between a tensor of shape (2, 3, 4) and a tensor of shape (1, 2, 3) will result in a tensor of shape (2, 6, 12).
To perform the normal kronecker product between matrices just pass rank-2 tensors.
To perform kronecker product between vectors first transfor them in rank-2 tensors of shape (1, -1)
Performing the kronecker product only along some legs means that along that leg it is an elementwise product and not a kronecker. For Example, if idxs=(0, 2) for the tensors of shapes (2, 3, 4) and (1, 3, 2) the output will be of shape (2, 3, 8).
- property linear_algebra_library
Specification of the linear algebra library used as string jax`.
- property links
Here, as well dimension of tensor along each dimension.
- mask_to_device(mask)[source]
Send a mask to the device where the tensor is. (right now only CPU –> GPU, CPU –> CPU).
- mask_to_host(mask)[source]
Send a mask to the host where we need it for symmetric tensors, e.g., degeneracies. Return as numpy.
- matmul(other: _AbstractQteaBaseTensor) QteaJaxTensor[source]
Concrete implementation of batched matrix multiplication for QteaJaxTensor.
Parameters
- other_AbstractQteaBaseTensor
Tensor with compatible shape for multiplication.
Returns
- QteaJaxTensor
New Tensor representing the (batched) matrix product.
- classmethod mpi_recv(from_, comm, tensor_backend)[source]
Send tensor via MPI.
Arguments
- from_integer
MPI process to receive tensor from.
comm : instance of MPI communicator to be used
tensor_backend : instance of
TensorBackend
- mpi_send(to_, comm)[source]
Send tensor via MPI.
Arguments
- tointeger
MPI process to send tensor to.
comm : instance of MPI communicator to be used
- property ndim
Rank of the tensor.
- norm_sqrt()[source]
Calculate the square root of the norm of the tensor, i.e., sqrt( <tensor|tensor>).
- permute_rows_cols_update(inds)[source]
Permute rows and columns of rank-2 tensor with inds. Inplace update.
- prepare_eig_api(conv_params)[source]
Return variables for eigsh.
Returns
- kwargsdict
Keyword arguments for eigs call. If initial guess can be passed, key “v0” is set with value None
LinearOperator : None
eigsh : None
- random_unitary(link)[source]
Generate a random unitary matrix via performing a SVD on a random tensor.
Arguments
- selfinstance of
QteaTensor Extract data type etc from this one here.
- linksame as returned by links property, here integer.
Dimension of the square, random unitary matrix.
- selfinstance of
- classmethod read(filehandle, dtype, device, base_tensor_cls, cmplx=True, order='F')[source]
Read a tensor from file via QteaTensor.
- scale_link(link_weights, link_idx, do_inverse=False)[source]
Scale tensor along one link at link_idx with weights.
Arguments
- link_weightsnp.ndarray
Scalar weights, e.g., singular values.
- link_idxint
Link which should be scaled.
- do_inversebool, optional
If True, scale with inverse instead of multiplying with link weights. Default to False
Returns
updated_link : instance of
QteaJaxTensorDetails
The inverse implementation handles zeros correctly which have been introduced due to padding. Therefore, scale_link should be used over passing 1 / link_weights to this function.
- scale_link_update(link_weights, link_idx, do_inverse=False)[source]
Scale tensor along one link at link_idx with weights (inplace update).
Arguments
- link_weightsnp.ndarray
Scalar weights, e.g., singular values.
- link_idxint
Link which should be scaled.
- do_inversebool, optional
If True, scale with inverse instead of multiplying with link weights. Default to False
Details
The inverse implementation handles zeros correctly which have been introduced due to padding. Therefore, scale_link should be used over passing 1 / link_weights to this function.
- set_diagonal_entry(position, value)[source]
Set the diagonal element in a rank-2 tensor (inplace update)
- set_matrix_entry(idx_row, idx_col, value)[source]
Set one element in a rank-2 tensor (inplace update)
- static set_seed(seed, devices=None)[source]
Set the seed for this tensor backend and the specified devices.
Arguments
- seedlist[int]
List of integers used as a seed; list has length 4.
- deviceslist[str] | None, optional
Can pass a list of devices via a string, e.g., to specify GPU by index. jax generates random numbers via a subkey which is passed to the function generating the random number and therefore should be device independent (we keep it for compatability). Default to None (set for all devices)
- set_submatrix(row_range, col_range, tensor)[source]
Set a submatrix of a rank-2 tensor for the given rows / cols.
- property shape
Dimension of tensor along each dimension.
- split_qr(legs_left, legs_right, perm_left=None, perm_right=None, is_q_link_outgoing=True, disable_streams=False)[source]
Split the tensor via a QR decomposition.
Parameters
- selfinstance of
QteaTensor Tensor upon which apply the QR
- legs_leftlist of int
Legs that will compose the rows of the matrix
- legs_rightlist of int
Legs that will compose the columns of the matrix
- perm_leftlist of int, optional
permutations of legs after the QR on left tensor
- perm_rightlist of int, optional
permutation of legs after the QR on right tensor
- disable_streamsboolean, optional
No effect here, but in general can disable streams to avoid nested generation of streams.
Returns
- tens_left: instance of
QteaTensor unitary tensor after the QR, i.e., Q.
- tens_right: instance of
QteaTensor upper triangular tensor after the QR, i.e., R
- selfinstance of
- split_qrte(tens_right, singvals_self, operator=None, conv_params=None, is_q_link_outgoing=True)[source]
Perform an Truncated ExpandedQR decomposition, generalizing the idea of https://arxiv.org/pdf/2212.09782.pdf for a general bond expansion given the isometry center of the network on tens_left. It should be rather general for three-legs tensors, and thus applicable with any tensor network ansatz. Notice that, however, you do not have full control on the approximation, since you know only a subset of the singular values truncated.
Parameters
- tens_left: xp.array
Left tensor
- tens_right: xp.array
Right tensor
- singvals_left: xp.array
Singular values array insisting on the link to the left of tens_left
- operator: xp.array or None
Operator to contract with the tensors. If None, no operator is contracted
Returns
- tens_left: ndarray
left tensor after the EQR
- tens_right: ndarray
right tensor after the EQR
- singvals: ndarray
singular values kept after the EQR
- singvals_cutted: ndarray
subset of thesingular values cutted after the EQR, normalized with the biggest singval
- split_svd(legs_left, legs_right, perm_left=None, perm_right=None, contract_singvals='N', conv_params=None, no_truncation=False, is_link_outgoing_left=True, disable_streams=False)[source]
Perform a truncated Singular Value Decomposition by first reshaping the tensor into a legs_left x legs_right matrix, and permuting the legs of the ouput tensors if needed. If the contract_singvals = (‘L’, ‘R’) it takes care of renormalizing the output tensors such that the norm of the MPS remains 1 even after a truncation.
Parameters
- selfinstance of
QteaTensor Tensor upon which apply the SVD
- legs_leftlist of int
Legs that will compose the rows of the matrix
- legs_rightlist of int
Legs that will compose the columns of the matrix
- perm_leftlist of int, optional
permutations of legs after the SVD on left tensor
- perm_rightlist of int, optional
permutation of legs after the SVD on right tensor
- contract_singvals: string, optional
- How to contract the singular values.
‘N’ : no contraction ‘L’ : to the left tensor ‘R’ : to the right tensor
- conv_params
TNConvergenceParameters, optional Convergence parameters to use in the procedure. If None is given, then use the default convergence parameters of the TN. Default to None.
- no_truncationboolean, optional
Allow to run without truncation Default to False (hence truncating by default)
- disable_streamsboolean, optional
No effect here, but in general can disable streams to avoid nested generation of streams.
Returns
- tens_left: instance of
QteaTensor left tensor after the SVD
- tens_right: instance of
QteaTensor right tensor after the SVD
- singvals: xp.ndarray
singular values kept after the SVD
- singvals_cut: xp.ndarray
singular values cut after the SVD, normalized with the biggest singval
- selfinstance of
- squeeze_update(position=None)[source]
Squeeze tensor at given position (inplace update).
Note that JAX’s squeeze returns a copy.
- stack_first_and_last_link(other)[source]
Stack first and last link of tensor targeting MPS addition.
- stack_link(other, link)[source]
Stack two tensors along a given link.
Arguments
- otherinstance of
QteaJaxTensor Links must match self up to the specified link.
- linkinteger
Stack along this link.
Returns
new_this : instance of :class:QteaJaxTensor`
- otherinstance of
- stream(disable_streams=False)[source]
Get the instance of a context which can be used to parallelize. Disabled for jax so far and always returning the nullcontext.
Parameters
- disable_streamsbool, optional
Allows to disable streams to avoid nested creation of streams. Globally, streams should be disabled via the set_streams_* function of the corresponding base tensor module (not available yet for jax). Default to False.
Returns
Context manager, i.e., always
nullcontext(AbstractContextManager)
- subtensor_along_link(link, lower, upper)[source]
Extract and return a subtensor select range (lower, upper) for one line.
- subtensor_along_link_inds(link, inds)[source]
Extract and return a subtensor via indices for one link.
Arguments
- linkint
Select only specific indices along this link (but all indices along any other link).
- indslist[int]
Indices to be selected and stored in the subtensor.
Returns
- subtensor
QteaJaxTensor Subtensor with selected indices.
Details
The numpy equivalent is
subtensor = tensor[:, :, inds, :]for a rank-4 tensor andlink=2.
- tensordot(other, contr_idx, disable_streams=False)[source]
Tensor contraction of two tensors along the given indices.
- to_dense(true_copy=False)[source]
Return dense tensor (if true_copy=False, same object may be returned).
- to_dense_singvals(s_vals, true_copy=False)[source]
Convert singular values to dense vector without symmetries.
- class qredtea.jaxapi.qteajaxtensor.DataMoverJax[source]
Data mover to move QteaJaxTensor between CPU and GPU.
- async_move(tensor, device, stream=None)[source]
Move the tensor tensor to the device device asynchronously with respect to the main computational stream
Parameters
- tensor_AbstractTensor
The tensor to be moved
- device: str
The device where to move the tensor
- streamstream-object
Stream to be used to move the data if a stream different from the data mover’s stream should be used. Default to None (use DataMover’s stream)
- property device_memory
Current memory occupied in the device